import cv2
import numpy as np
from ultralytics import YOLO                # 导入Ultralytics的YOLOv8模型库
from typing import Tuple, List, Dict, Any, Optional
from pathlib import Path
import torch                                # 用于判断CUDA（GPU）是否可用
import os


class YOLOv8Detector:
    def __init__(self, config_loader):
        # 读取模型相关配置（如模型路径、设备等）
        self.config = config_loader.get_model_config()
        # 读取应用相关配置（如检测类别、置信度等）
        self.app_config = config_loader.get_app_config()
        self.model_person = None               # YOLO模型对象，人员模型
        self.model_fire = None               # YOLO模型对象，fire模型
        self.is_initialized = False     # 标记模型是否加载成功
        self._load_model()              # 初始化时自动加载模型

    def _load_model(self):
        """加载YOLOv8模型（支持自动切换CPU/GPU）"""
        try:
             # 获取项目根目录（main.py所在目录）
            project_root = os.path.dirname(os.path.abspath(__file__))
            project_root = os.path.abspath(os.path.join(project_root, "../../"))  # 根据你的src目录深度调整
            
            # 人员模型 ------------------------------------------------------------
            person_cfg_md = self.config['model']['person']       
            device = person_cfg_md['device']         # 获取期望的运行设备（'cuda'或'cpu'）

            # 检查CUDA（GPU）是否可用，如果不可用则自动切换到CPU
            if device == 'cuda' and not torch.cuda.is_available():
                print("CUDA不可用，切换到CPU")
                device = 'cpu'

            self.model_person = YOLO(person_cfg_md['path'])           # 加载YOLOv8模型文件
            self.model_person.to(device)                   # 将模型移动到指定设备（GPU或CPU）
            # if self.config['optimization'].get('half', False):  # 检查是否启用半精度
            #     self.model_person.half()  # 转换为半精度

            # 火焰模型 ------------------------------------------------------------
            fire_cfg_md = self.config['model']['fire']  # 获取火焰模型配置
            fire_path = fire_cfg_md['path']

            # TODO: 如果想对路径不行就改用绝对路径，解除这一段的注释
            # if not os.path.isabs(fire_path):
            #     fire_path = os.path.join(project_root, fire_path)
            #     print(f"火焰模型路径: {fire_path}")

            self.model_fire = YOLO(fire_path)  # 加载火焰检测模型
            self.model_fire.to(device)

            self.is_initialized = True              # 标记模型已加载
            print(f"模型加载成功: person--{person_cfg_md['path']}  fire--{fire_cfg_md['path']}")      # 打印加载成功信息
        except Exception as e:
            print(f"模型加载失败: {e}")               # 打印加载失败信息
            self.is_initialized = False             # 标记模型未加载

    def detect(self, frame: np.ndarray) -> Tuple[bool, List[Dict[str, Any]], bool]:
        """
         对输入帧进行目标检测，返回检测结果和危险标记

        返回:
            - is_detected: 是否成功执行检测
            - detections: 检测结果列表
            - has_danger: 是否检测到危险物品
        """
        if not self.is_initialized:     # 如果模型未初始化，直接返回检测失败
            return False, [], False

        try:
            detections = []             # 初始化检测结果列表
            has_danger = False          # 初始化危险标记

            # 执行检测yolov8模型推理,(person)
            results_person = self.model_person(
                frame,
                imgsz=self.config['model']['person']['imgsz'],                        # 输入图像大小
                conf=self.app_config['detection']['confidence_threshold'],  # 置信度阈值
                max_det=self.config['optimization']['max_det'],             # 最大检测数量
                verbose=False                                               # 不输出详细日志
            )

            # 获取配置中定义的检测类别信息
            detection_config = self.app_config['detection']['classes']
            # 解析检测结果
            for r in results_person:
                boxes = r.boxes                         # 获取检测框
                if boxes is not None:                   # 如果检测框不为空
                    for box in boxes:                   # 遍历每个检测框
                        cls_id = int(box.cls)           # 获取类别ID
                        conf = float(box.conf)          # 获取置信度
                        xyxy = box.xyxy[0].tolist()     # 获取边界框坐标


                        # 人
                        if cls_id == detection_config['person']['id']:
                            detections.append({
                                'class': 'person',          # 类别名称
                                'confidence': conf,         # 置信度
                                'bbox': xyxy,               # 边界框坐标
                                'color': detection_config['person']['color'],           # 颜色
                                'is_danger': detection_config['person']['is_danger']    # 是否危险
                            })

                        # 杯子
                        elif cls_id == detection_config['cup']['id']:
                            detections.append({
                                'class': 'cup',
                                'confidence': conf,
                                'bbox': xyxy,
                                'color': detection_config['cup']['color'],
                                'is_danger': detection_config['cup']['is_danger']
                            })
                            # has_danger = True   # 标记检测到危险物品

                        # 注意：火焰检测可能需要自定义模型
                        # 这里预留接口
            # 2. 火焰检测
            results_fire = self.model_fire(
                frame,
                imgsz=self.config['model']['fire']['imgsz'],                # 设置输入图像尺寸
                conf=self.app_config['detection']['confidence_threshold'],  # 设置置信度阈值
                max_det=self.config['optimization']['max_det'],             # 设置最大检测数
                verbose=False  # 关闭详细输出
            )
            fire_id = detection_config['fire']['id']    # 获取火焰类别ID
            for r in results_fire:                      # 遍历检测结果
                boxes = r.boxes                         # 获取边界框
                if boxes is not None:                   # 如果有检测结果
                    for box in boxes:                   # 遍历每个边界框
                        cls_id = int(box.cls)           # 获取类别ID
                        conf = float(box.conf)          # 获取置信度
                        xyxy = box.xyxy[0].tolist()     # 获取边界框坐标
                        if cls_id == fire_id:           # 如果是火焰
                            detections.append({
                                'class': 'fire',        # 类别名称
                                'confidence': conf,     # 置信度
                                'bbox': xyxy,           # 边界框坐标
                                'color': detection_config['fire']['color'],         # 显示颜色
                                'is_danger': detection_config['fire']['is_danger']  # 是否危险
                            })
                            has_danger = True           # 标记存在危险

            # 返回检测结果：是否成功、检测到的目标列表、是否存在危险
            return True, detections, has_danger  

        except Exception as e:
            print(f"检测错误: {e}")
            return False, [], False
